import pandas as pd
from simple_visualization import SimpleTradingVisualizer

class TradingSimulator:
    def __init__(self, initial_capital=1000000, transaction_cost=0.001,
                 risk_per_trade=0.02, min_trade_interval=60):
        self.initial_capital = initial_capital
        self.transaction_cost = transaction_cost
        self.risk_per_trade = risk_per_trade
        self.min_trade_interval = min_trade_interval
        self.reset()

    def reset(self):
        # 重置模拟器状态
        self.capital = self.initial_capital
        self.equity_curve = []
        self.current_position = None
        self.trade_history = []
        self.last_trade_time = None
        self.signals = []

    def calculate_atr(self, data, window=14):
        # 计算平均真实波幅(ATR)
        data = data.copy()
        high = data['high']
        low = data['low']
        close = data['close']

        # 计算真实波幅(TR)
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))

        tr = pd.DataFrame({'tr1': tr1, 'tr2': tr2, 'tr3': tr3}).max(axis=1)
        atr = tr.rolling(window=window).mean()

        return atr

    def calculate_position_size(self, entry_price, stop_loss_price, volatility, confidence=1.0):
        # 基于置信度的仓位管理
        risk_per_share = abs(entry_price - stop_loss_price)
        if risk_per_share == 0:
            return 0

        # 置信度调整：高质量信号使用更大仓位
        vol_adjustment = min(1.0, 0.02 / (volatility + 1e-6))
        confidence_adjustment = min(1.5, max(0.5, confidence))  # 置信度在0.5-1.5之间调整
        max_risk_amount = self.capital * self.risk_per_trade * vol_adjustment * confidence_adjustment
        position_size = max_risk_amount / risk_per_share

        max_position_value = self.capital * 0.8
        calculated_value = position_size * entry_price
        if calculated_value > max_position_value:
            position_size = max_position_value / entry_price

        return int(position_size)

    def calculate_stop_loss(self, entry_price, signal, atr_value, data, current_idx):
        # 改进的动态止损策略
        if atr_value is None:
            atr_value = entry_price * 0.015

        # 基于支撑阻力位的止损
        if current_idx >= 20:  # 确保有足够的历史数据
            recent_low = data['low'].iloc[current_idx - 20:current_idx].min()
            recent_high = data['high'].iloc[current_idx - 20:current_idx].max()

            if signal == 2:  # 做多
                # 将止损设在近期低点下方一定距离
                support_stop = recent_low - atr_value * 0.5
                atr_stop = entry_price - atr_value * 1.5
                return max(0.01, min(support_stop, atr_stop))
            else:  # 做空
                # 将止损设在近期高点上方的距离
                resistance_stop = recent_high + atr_value * 0.5
                atr_stop = entry_price + atr_value * 1.5
                return resistance_stop if atr_stop > resistance_stop else atr_stop

        # 默认ATR止损
        if signal == 2:
            return max(0.01, entry_price - atr_value * 1.5)
        else:
            return entry_price + atr_value * 1.5

    def calculate_take_profit(self, entry_price, stop_loss, signal):
        # 动态止盈，风险回报比1:3
        risk_amount = abs(entry_price - stop_loss)

        if signal == 2:  # 做多
            return entry_price + (risk_amount * 2.0)  # 风险回报比1:2
        else:  # 做空
            return max(0.01, entry_price - (risk_amount * 2.0))

    def confirm_volume(self, data, current_idx):
        # 成交量确认
        if current_idx < 5:
            return False

        # 当前成交量高于近期平均
        current_volume = data['vol'].iloc[current_idx]
        avg_volume = data['vol'].iloc[current_idx - 5:current_idx].mean()

        if current_volume < avg_volume * 1.2:  # 需要放量20%以上
            return False

        return True

    def confirm_trend(self, data, current_idx, signal):
        # 多指标趋势确认
        if current_idx < 2:  # 减少到2个周期
            return False

        # MA趋势确认
        ma_short = data['MA_short'].iloc[current_idx]
        ma_medium = data['MA_medium'].iloc[current_idx]

        if signal == 2:  # 做多只需要短期>中期
            if not (ma_short > ma_medium):
                return False
        else:  # 做空只需要短期<中期
            if not (ma_short < ma_medium):
                return False

        # MACD确认
        macd = data['MACD'].iloc[current_idx]

        if signal == 2 and macd < 0:  # 做多时MACD只要为正
            return False
        if signal == 0 and macd > 0:  # 做空时MACD只要为负
            return False

        return True

    def is_trend_aligned(self, data, current_idx, signal):
        # 检查是否顺势交易
        if current_idx < 20:
            return True

        # 检查长期趋势
        ma_long = data['MA_long'].iloc[current_idx]
        ma_long_prev = data['MA_long'].iloc[current_idx - 1]

        long_trend = 2 if ma_long > ma_long_prev else 0  # 2=上升趋势, 0=下降趋势

        # 只做顺势交易：上升趋势只做多，下降趋势只做空
        if long_trend == 2 and signal != 2:
            return False
        if long_trend == 0 and signal != 0:
            return False

        return True

    def check_support_resistance(self, data, current_idx, signal, current_price):
        # 检查支撑阻力位
        if current_idx < 30:
            return True

        recent_low = data['low'].iloc[current_idx - 20:current_idx].min()
        recent_high = data['high'].iloc[current_idx - 20:current_idx].max()

        if signal == 2:  # 做多
            # 放松到支撑位上5%
            support_level = recent_low * 1.05
            return current_price <= support_level
        else:  # 做空
            # 放松到阻力位下5%
            resistance_level = recent_high * 0.95
            return current_price >= resistance_level

    def check_trend_strength(self, data, current_idx, signal):
        # 趋势强度检查
        if current_idx < 10:
            return True

        # ADX条件检查
        if 'ADX' in data.columns:
            adx_value = data['ADX'].iloc[current_idx]
            if adx_value < 20:  # 从25降低到20
                return False

        return True

    def check_price_position(self, data, current_idx):
        # 检查价格位置（避免在布林带极端位置入场）
        if 'BB_upper' in data.columns and 'BB_lower' in data.columns:
            current_close = data['close'].iloc[current_idx]
            bb_upper = data['BB_upper'].iloc[current_idx]
            bb_lower = data['BB_lower'].iloc[current_idx]
            bb_middle = (bb_upper + bb_lower) / 2

            # 避免在布林带上下轨附近入场
            if abs(current_close - bb_upper) / (bb_upper - bb_middle) < 0.1:
                return False
            if abs(current_close - bb_lower) / (bb_middle - bb_lower) < 0.1:
                return False

        return True

    def should_enter_trade(self, date, pred, current_price, current_atr, test_data, i):
        # 交易信号过滤
        if not self.is_trend_aligned(test_data, i, pred):
            return False

        if self.current_position is not None:
            return False

        if pred == 1:  # 观望信号
            return False

        # 交易间隔检查
        if self.last_trade_time is not None:
            time_diff = (date - self.last_trade_time).total_seconds() / 3600
            if time_diff < 4:  # 4小时内不重复交易
                return False

        # 波动率过滤 - 放松到3%
        if current_atr is not None and current_atr / current_price > 0.03:
            return False

        # 趋势确认 - 放松条件
        if not self.confirm_trend(test_data, i, pred):
            return False

        # 波动率过滤 - 避免低波动环境
        if current_atr is not None and current_atr / current_price < 0.005:  # 波动率低于0.5%不交易
            return False

        # 支撑阻力确认
        if not self.check_support_resistance(test_data, i, pred, current_price):
            return False
        # 增加趋势强度过滤
        if not self.check_trend_strength(test_data, i, pred):
            return False

        # 增加价格位置过滤（避免在极端位置入场）
        if not self.check_price_position(test_data, i):
            return False

        return True

    def enter_position_with_params(self, date, signal, price, stop_loss, take_profit, position_size):
        # 带参数的开仓方法
        if self.current_position is not None:
            return False

        trade_cost = position_size * price * self.transaction_cost

        if trade_cost + (position_size * price) > self.capital:
            return False

        self.current_position = {
            'entry_date': date,
            'signal': signal,
            'entry_price': price,
            'position_size': position_size,
            'stop_loss': stop_loss,
            'take_profit': take_profit,
            'entry_cost': trade_cost
        }

        self.capital -= trade_cost

        self.signals.append({
            'date': date,
            'type': 'entry',
            'signal': signal,
            'price': price,
            'action': 'long' if signal == 2 else 'short'
        })

        print(f"{date.strftime('%Y-%m-%d %H:%M')}: 开仓 {'做多' if signal == 2 else '做空'} "
              f"价格={price:.2f}, 数量={position_size}, 止损={stop_loss:.2f}, 止盈={take_profit:.2f}")
        return True

    def exit_position(self, date, exit_price, exit_reason):
        # 平仓
        if self.current_position is None:
            return

        position = self.current_position
        entry_value = position['position_size'] * position['entry_price']
        exit_value = position['position_size'] * exit_price

        if position['signal'] == 2:  # 做多
            pnl = exit_value - entry_value
        else:  # 做空
            pnl = entry_value - exit_value

        # 扣除交易成本
        exit_cost = position['position_size'] * exit_price * self.transaction_cost
        net_pnl = pnl - position['entry_cost'] - exit_cost

        # 更新资金
        self.capital += net_pnl

        # 记录交易
        trade_record = {
            'entry_date': position['entry_date'],
            'exit_date': date,
            'signal': position['signal'],
            'entry_price': position['entry_price'],
            'exit_price': exit_price,
            'position_size': position['position_size'],
            'pnl': net_pnl,
            'return_pct': (net_pnl / entry_value) * 100 if entry_value > 0 else 0,
            'exit_reason': exit_reason,
            'trade_type': 'long' if position['signal'] == 2 else 'short'
        }

        self.trade_history.append(trade_record)

        # 记录平仓信号
        self.signals.append({
            'date': date,
            'type': 'exit',
            'signal': position['signal'],
            'price': exit_price,
            'action': 'exit_long' if position['signal'] == 2 else 'exit_short'
        })

        self.current_position = None

        print(f"{date.strftime('%Y-%m-%d %H:%M')}: 平仓 {exit_reason}, "
              f"价格={exit_price:.2f}, 盈亏={net_pnl:.2f} ({trade_record['return_pct']:.2f}%)")

    def check_exit_conditions(self, date, current_price, current_low, current_high):
        # 检查退出条件
        if self.current_position is None:
            return False, None, None

        position = self.current_position

        # 移动止损逻辑
        if position['signal'] == 2:  # 做多
            # 价格上涨时提高止损
            new_stop_loss = max(position['stop_loss'], current_price * 0.99)  # 移动止损到99%
            position['stop_loss'] = new_stop_loss
        else:  # 做空
            # 价格下跌时降低止损
            new_stop_loss = min(position['stop_loss'], current_price * 1.01)  # 移动止损到101%
            position['stop_loss'] = new_stop_loss

        # 检查止损
        if position['signal'] == 2:  # 做多
            if current_low <= position['stop_loss']:
                return True, position['stop_loss'], '止损'
            elif current_high >= position['take_profit']:
                return True, position['take_profit'], '止盈'
        else:  # 做空
            if current_high >= position['stop_loss']:
                return True, position['stop_loss'], '止损'
            elif current_low <= position['take_profit']:
                return True, position['take_profit'], '止盈'

        return False, None, None

    def run_simulation(self, test_data, predictions, lookback=30):
        # 运行交易模拟
        self.reset()
        test_data = test_data.copy()
        test_data['atr'] = self.calculate_atr(test_data)

        test_start_idx = lookback
        valid_predictions = min(len(predictions), len(test_data) - test_start_idx)

        test_dates = test_data.index[test_start_idx:test_start_idx + valid_predictions]
        test_prices = test_data['close'].iloc[test_start_idx:test_start_idx + valid_predictions]
        test_high = test_data['high'].iloc[test_start_idx:test_start_idx + valid_predictions]
        test_low = test_data['low'].iloc[test_start_idx:test_start_idx + valid_predictions]
        test_atr = test_data['atr'].iloc[test_start_idx:test_start_idx + valid_predictions]

        signal_count = 0
        trade_count = 0

        for i in range(valid_predictions):
            date = test_dates[i]
            pred = predictions[i]
            current_price = test_prices.iloc[i]
            current_high = test_high.iloc[i]
            current_low = test_low.iloc[i]
            current_atr = test_atr.iloc[i] if not pd.isna(test_atr.iloc[i]) else None

            # 检查退出条件
            should_exit, exit_price, exit_reason = self.check_exit_conditions(
                date, current_price, current_low, current_high
            )

            if should_exit:
                self.exit_position(date, exit_price, exit_reason)
                trade_count += 1
                self.last_trade_time = date

            # 检查开仓条件 - 使用严格的过滤条件
            if self.current_position is None and pred != 1:
                if self.should_enter_trade(date, pred, current_price, current_atr, test_data, i):
                    signal_count += 1
                    if signal_count <= 10:
                        print(f"{date.strftime('%Y-%m-%d %H:%M')}: 发现高质量交易信号 {pred}，价格 {current_price:.2f}")

                    # 动态计算止损止盈
                    stop_loss = self.calculate_stop_loss(current_price, pred, current_atr, test_data,
                                                         test_start_idx + i)
                    take_profit = self.calculate_take_profit(current_price, stop_loss, pred)

                    # 基于波动率的仓位计算
                    volatility = current_atr / current_price if current_atr else 0.02
                    position_size = self.calculate_position_size(current_price, stop_loss, volatility)

                    if position_size > 0:
                        self.enter_position_with_params(date, pred, current_price, stop_loss, take_profit,
                                                        position_size)
                        if self.current_position is not None:
                            trade_count += 1
                            self.last_trade_time = date

            # 检查持仓时间过长强制平仓
            if self.current_position is not None:
                holding_period = (date - self.current_position['entry_date']).total_seconds() / 3600  # 小时
                if holding_period > 48:  # 最多持有48小时
                    self.exit_position(date, current_price, '持仓时间过长')
                    trade_count += 1
                    self.last_trade_time = date

            # 记录每日权益
            daily_equity = self.calculate_daily_equity(current_price)
            self.equity_curve.append({
                'date': date,
                'equity': daily_equity,
                'price': current_price
            })

        # 强制平仓
        if self.current_position is not None:
            self.exit_position(test_dates[-1], test_prices.iloc[-1], '强制平仓')
            trade_count += 1

        print(f"模拟完成！总共执行 {trade_count} 笔交易")

    def calculate_daily_equity(self, current_price):
        # 计算每日权益
        if self.current_position is None:
            return self.capital

        position = self.current_position
        position_value = position['position_size'] * current_price
        entry_value = position['position_size'] * position['entry_price']

        if position['signal'] == 2:  # 做多
            unrealized_pnl = position_value - entry_value
        else:  # 做空
            unrealized_pnl = entry_value - position_value

        return self.capital + unrealized_pnl - position['entry_cost']

    def generate_simple_visualizations(self, price_data=None):
        visualizer = SimpleTradingVisualizer(
            self.equity_curve,
            self.trade_history,
            self.signals,
            price_data
        )

        visualizer.plot_all()

    def generate_performance_report(self):
        # 生成绩效报告
        if not self.trade_history:
            print("没有交易记录")
            return None, None

        trades_df = pd.DataFrame(self.trade_history)
        equity_df = pd.DataFrame(self.equity_curve)

        if equity_df.empty:
            print("没有权益数据")
            return trades_df, equity_df

        # 基本统计
        total_trades = len(trades_df)
        winning_trades = len(trades_df[trades_df['pnl'] > 0])
        losing_trades = len(trades_df[trades_df['pnl'] <= 0])
        win_rate = winning_trades / total_trades if total_trades > 0 else 0

        total_pnl = trades_df['pnl'].sum()
        avg_win = trades_df[trades_df['pnl'] > 0]['pnl'].mean() if winning_trades > 0 else 0
        avg_loss = trades_df[trades_df['pnl'] <= 0]['pnl'].mean() if losing_trades > 0 else 0
        profit_factor = abs(avg_win * winning_trades) / abs(avg_loss * losing_trades) if losing_trades > 0 else float(
            'inf')

        # 最大回撤
        equity_df['peak'] = equity_df['equity'].cummax()
        equity_df['drawdown'] = (equity_df['equity'] - equity_df['peak']) / equity_df['peak']
        max_drawdown = equity_df['drawdown'].min() * 100 if not equity_df.empty else 0

        # 年化收益
        if len(equity_df) > 1:
            days = (equity_df['date'].iloc[-1] - equity_df['date'].iloc[0]).days
            if days > 0:
                total_return = equity_df['equity'].iloc[-1] / self.initial_capital - 1
                annual_return = (1 + total_return) ** (365 / days) - 1
            else:
                annual_return = 0
        else:
            annual_return = 0

        print("=" * 60)
        print("交易绩效报告")
        print("=" * 60)
        print(f"初始资金: {self.initial_capital:,.2f}")
        print(f"最终资金: {self.capital:,.2f}")
        print(f"总盈亏: {total_pnl:,.2f}")
        print(f"总收益率: {(self.capital / self.initial_capital - 1) * 100:.2f}%")
        print(f"年化收益率: {annual_return * 100:.2f}%")
        print(f"总交易次数: {total_trades}")
        print(f"胜率: {win_rate * 100:.2f}%")
        print(f"平均盈利: {avg_win:,.2f}")
        print(f"平均亏损: {avg_loss:,.2f}")
        print(f"盈亏比: {profit_factor:.2f}")
        print(f"最大回撤: {max_drawdown:.2f}%")

        # 详细的交易统计
        if not trades_df.empty:
            print("\n详细交易统计:")
            print(f"做多交易: {len(trades_df[trades_df['signal'] == 2])}")
            print(f"做空交易: {len(trades_df[trades_df['signal'] == 0])}")
            print(f"最大单笔盈利: {trades_df['pnl'].max():.2f}")
            print(f"最大单笔亏损: {trades_df['pnl'].min():.2f}")

        print("=" * 60)